import json
import requests
from loguru import logger
import os
import hashlib
import gzip
import tarfile
import shutil
from core.base_class import BaseClass
import urllib3
urllib3.disable_warnings()


class ImageDownloader(BaseClass):
    def __init__(self):
        BaseClass.__init__(self)
        self.images = json.loads(self.config["要下载的镜像"])
        self.proxy = self.config["代理地址"]

    def get_image_info(self, image: str) -> (str, str, str, str):
        registry = self.proxy if self.proxy else 'registry-1.docker.io'
        repo = 'library'
        tag = 'latest'
        image_parts = image.split('/')
        try:
            img, tag = image_parts[-1].split('@')
        except ValueError:
            try:
                img, tag = image_parts[-1].split(':')
            except ValueError:
                img = image_parts[-1]

        if len(image_parts) > 1 and ('.' in image_parts[0] or ':' in image_parts[0]):
            registry = image_parts[0]
            repo = '/'.join(image_parts[1:-1])
        else:
            if len(image_parts[:-1]) != 0:
                repo = '/'.join(image_parts[:-1])

        repository = '{}/{}'.format(repo, img)
        return registry, repository, img, tag

    @staticmethod
    def get_auth_header(registry, repo) -> dict[str, str]:
        auth_url = 'https://auth.docker.io/token'
        reg_service = 'registry.docker.io'
        resp = requests.get('https://{}/v2/'.format(registry), verify=False)
        if resp.status_code == 401:
            auth_url = resp.headers['WWW-Authenticate'].split('"')[1]
            try:
                reg_service = resp.headers['WWW-Authenticate'].split('"')[3]
            except IndexError:
                reg_service = ""

        resp = requests.get(f'{auth_url}?service={reg_service}&scope=repository:{repo}:pull', verify=False)
        access_token = resp.json()['token']
        return {
            'Authorization': 'Bearer ' + access_token,
            'Accept': "application/vnd.docker.distribution.manifest.v2+json"
        }

    @staticmethod
    def get_manifests(registry, repo, tag, headers) -> (str | None, str | None):
        resp = requests.get(
            'https://{}/v2/{}/manifests/{}'.format(registry, repo, tag),
            headers=headers, verify=False
        )
        if resp.status_code != 200:
            logger.error('Cannot fetch manifest for {} [HTTP {}]'.format(repo, resp.status_code))
            return None, None
        response = resp.json()
        logger.debug(response)
        if not response.__contains__("layers"):
            logger.error("manifests没有layers, 不支持OCI格式")
            exit(0)
        return response['layers'], response['config']['digest']

    @staticmethod
    def get_blobs(registry, repo, digest, headers):
        resp = requests.get(
            'https://{}/v2/{}/blobs/{}'.format(registry, repo, digest),
            headers=headers, verify=False
        )
        return resp.content

    @staticmethod
    def download_layers(layers, img_dir, registry, repo, headers, content, blobs):
        empty_json = '{"created":"1970-01-01T00:00:00Z","container_config":{"Hostname":"","Domainname":"","User":"",'\
                      '"AttachStdin":false, "AttachStdout":false,"AttachStderr":false,"Tty":false,"OpenStdin":false, '\
                      '"StdinOnce":false,"Env":null,"Cmd":null,"Image":"", "Volumes":null,"WorkingDir":"",'\
                      '"Entrypoint":null,"OnBuild":null,"Labels":null}}'
        parent_id = ""
        fake_layer_id = ""
        for layer in layers:
            layer_digest = layer['digest']
            # todo: Creating fake layer ID. Don't know how Docker generates it
            fake_layer_id = hashlib.sha256((parent_id + '\n' + layer_digest + '\n').encode('utf-8')).hexdigest()
            layer_dir = img_dir + '/' + fake_layer_id
            os.mkdir(layer_dir)

            # Creating VERSION file
            with open(layer_dir + '/VERSION', 'w') as fd:
                fd.write('1.0')

            logger.debug(f"Downloading layer[{layer_digest[7:19]}] ...")
            resp = requests.get('https://{}/v2/{}/blobs/{}'.format(registry, repo, layer_digest),
                                headers=headers, stream=True, verify=False)
            if resp.status_code != 200:  # When the layer is located at a custom URL
                resp = requests.get(layer['urls'][0], headers=headers, stream=True, verify=False)
                if resp.status_code != 200:
                    logger.error(f'Cannot download {repo} layer [HTTP {resp.status_code}]')
                    return

            resp.raise_for_status()
            unit = int(resp.headers['Content-Length']) / 50
            acc = 0
            nb_traits = 0
            with open(layer_dir + '/layer_gzip.tar', "wb") as file:
                for chunk in resp.iter_content(chunk_size=8192):
                    if chunk:
                        file.write(chunk)
                        acc = acc + 8192
                        if acc > unit:
                            nb_traits = nb_traits + 1
                            acc = 0
            # 提取
            with open(layer_dir + '/layer.tar', "wb") as file:  # Decompress gzip response
                unzip_layer = gzip.open(layer_dir + '/layer_gzip.tar', 'rb')
                shutil.copyfileobj(unzip_layer, file)
                unzip_layer.close()
            os.remove(layer_dir + '/layer_gzip.tar')

            content[0]['Layers'].append(fake_layer_id + '/layer.tar')

            with open(layer_dir + '/json', 'w') as file:
                # last layer = config manifest - history - rootfs
                if layers[-1]['digest'] == layer['digest']:
                    # todo: json.loads() automatically converts to unicode, thus decoding values whereas Docker doesn't
                    json_obj = json.loads(blobs)
                    del json_obj['history']
                    try:
                        del json_obj['rootfs']
                    except:  # Because Microsoft loves case insensitiveness
                        del json_obj['rootfS']
                else:  # other layers json are empty
                    json_obj = json.loads(empty_json)
                json_obj['id'] = fake_layer_id
                if parent_id:
                    json_obj['parent'] = parent_id
                parent_id = json_obj['id']
                file.write(json.dumps(json_obj))
        return fake_layer_id

    def start(self):
        for image in self.images:
            registry, repo, image_name, tag = self.get_image_info(image)
            # logger.debug(f"{registry}-{repo}-{image_name}-{tag}")
            auth_header = self.get_auth_header(registry, repo)
            logger.debug(auth_header)
            # 获取元数据
            layers, digest = self.get_manifests(registry, repo, tag, auth_header)
            if layers is None:
                continue
            # 创建临时目录，用于存放数据
            img_dir = f'tmp_{image_name}@{tag}'
            if os.path.exists(img_dir):
                shutil.rmtree(img_dir)
            os.mkdir(img_dir)
            #
            blobs = self.get_blobs(registry, repo, digest, auth_header)
            with open('{}/{}.json'.format(img_dir, digest[7:]), 'wb') as fd:
                fd.write(blobs)

            content = [{
                'Config': digest[7:] + '.json',
                'RepoTags': [f'{registry}/{repo}:{tag}'],
                'Layers': []
            }]
            # 下载layers
            fake_layer_id = self.download_layers(layers, img_dir, registry, repo, auth_header, content, blobs)

            with open(img_dir + '/manifest.json', 'w') as file:
                file.write(json.dumps(content))

            content = {f'{registry}/{repo}': {tag: fake_layer_id}}

            with open(img_dir + '/repositories', 'w') as file:
                file.write(json.dumps(content))

            # 生成镜像tar文件
            docker_tar = repo.replace('/', '_') + '.tar'

            tar = tarfile.open(docker_tar, "w")
            tar.add(img_dir, arcname=os.path.sep)
            tar.close()
            shutil.rmtree(img_dir)
            logger.success(f"create [{repo}] success!")


if __name__ == '__main__':
    downloader = ImageDownloader()
    downloader.start()
